# -*- coding: utf-8 -*-  
'''
GPT2.0 评价函数

@author: luoyi
Created on 2021年4月7日
'''
import tensorflow as tf

import utils.conf as conf
import utils.logger_factory as logf


#    预训练评价函数
class Gpt2PreMetric(tf.keras.metrics.Metric):
    def __init__(self,
                 name='Gpt2PreMetric',
                 sentence_maxlen=conf.GPT2.get_pre_training_sentence_maxlen(),
                 **kwargs):
        super(Gpt2PreMetric, self).__init__(name=name, **kwargs)
        
        self._sentence_maxlen = sentence_maxlen
        
        self._acc = self.add_weight('pre_acc', initializer='zero', dtype=tf.float64)
        pass
    
    def update_state(self, y_true, y_pred, sample_weight=None):
        '''
            @param y_pred: Tensor(batch_size, sentence_maxlen, vocab_size)
            @param y_true: Tensor(batch_size, sentence_maxlen)
        '''
        y_true = tf.cast(y_true, dtype=tf.int64)
        
        #    通过y_true中 > 0的索引拿到每个实际需要预测的词
        idx = tf.where(y_true > 0)                      #    Tensor(sum(每个batch中实际的词数量), 3)
        count = tf.math.count_nonzero(y_true, axis=-1)  #    RaggedTendor(batch_size, None(每个batch中实际有效词数量))
        #    y_true中真实的词编码
        y_true = tf.gather_nd(y_true, indices=idx)      #    Tensor(sum(每个batch中实际的词数量), )
           
        #    预测的词编码
        y_pred = tf.gather_nd(y_pred, indices=idx)      #    Tensor(sum(每个batch中实际的词数量), vocab_size)
        y_pred = tf.math.argmax(y_pred, axis=-1)        #    Tensor(sum(每个batch中实际的词数量), )
        
        #    每一个词的预测结果 True:正确，False:错误
        equal_res = tf.equal(y_true, y_pred)            #    Tensor(sum(每个batch中实际的词数量), )
#         equal_res = tf.RaggedTensor.from_row_lengths(equal_res, row_lengths=count)      #    RaggedTendor(batch_size, None(每个batch中预测正确的词数量))
        
        P = tf.math.reduce_sum(count)                   #    总需要预测的词数量
        T = tf.math.count_nonzero(equal_res)            #    总预测正确的词数量
        
        acc = T / P
        self._acc.assign(acc)
        
        tf.print('------------------------------------------------', output_stream=logf.get_logger_filepath('gpt2_metrics'))
        tf.print('acc:', acc, ' T:', T, ' P:', P, output_stream=logf.get_logger_filepath('gpt2_metrics'))
        tf.print(y_pred, output_stream=logf.get_logger_filepath('gpt2_metrics'))
        pass
    def reset_states(self):
        self._acc.assign(0.)
        pass
    def result(self):
        return self._acc
    pass
